Infer dtype of Series in more cases #766

twoertwein · 2023-08-17T23:37:16Z

Closes Create pd.Series initializer overload to propagate the generic type of an input container #710 (Replace xxxx with the Github issue number)
Tests added: Please use assert_type() to assert the type of any return value

pandas-stubs/core/frame.pyi

pandas-stubs/core/series.pyi

tests/test_series.py

twoertwein · 2023-08-17T23:54:01Z

I don't know why both mypy and pyright think that Series[int].xs is int!

Dr-Irv

I saw this is in draft form, but I think my comments would need to be addressed anyway

pandas-stubs/core/frame.pyi

pandas-stubs/core/series.pyi

tests/test_series.py

Dr-Irv · 2023-08-18T12:35:03Z

tests/test_series.py

    assert assert_type(s.replace(1, 2, inplace=True), None) is None


 def test_cat_accessor() -> None:
    # GH 43
-    s = pd.Series(pd.Categorical(["a", "b", "a"], categories=["a", "b"]))
+    s: pd.Series[str] = pd.Series(
+        pd.Categorical(["a", "b", "a"], categories=["a", "b"])


Maybe we should have a CategoricalSeries ?

I think we might need to reconsider S1 at some point (not in this PR). I would slightly prefer S1 = TypeVar("S1") as any type can be in a Series.

Yes, but I'd rather encourage "typical" types to be in a Series, and let people do cast or ignore if they are putting other types in a Series.

We could do a CategoricalSeries as a future PR

tests/test_series.py

Dr-Irv

It all looks good. Some open comments to address. I think this is the list:

Where you changed the construction of the Series to shorten them, I'd rather leave them as is to have less changes to the tests.
For Series.update(), need to surround that ignore with a TYPE_CHECKING_INVALID_USAGE pattern
Use np.integer rather than np.intp
For strings passed into Series.agg(), maybe we just need to have an overload for np.mean and "mean" to say they return Series[float] (as opposed to what I wrote elsewhere about returning Any) and then you can remove the ignore in the tests.

twoertwein · 2023-08-18T17:39:57Z

I believe I fixed everything but agg/aggregate. will look into whether we can do something like:

special case "mean"
any string return S1
for callables: return the actual return type

Dr-Irv · 2023-08-18T18:42:43Z

pandas-stubs/core/series.pyi

@@ -834,7 +859,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
        axis: AxisIndex = ...,
        *args,
        **kwargs,
-    ) -> Series[S1]: ...
+    ) -> Self: ...


Maybe make this Series since we really don't know the type of the Series - could be mix of int and float, or one, or the other

Thank you, that made it work!

Dr-Irv

Pretty close. Just that one case where there is a mix of funcs on the Series.agg() call. Maybe the solution is that if a list is passed, the result is Series instead of Series[S1] . See comment

Dr-Irv

thanks @twoertwein

Infer dtype of Series in more cases

66ece55